sleap_nn/config/training_job_config.py (3)

154-162: Missing docstring for new function

Unlike other functions in this file, load_sleap_config lacks a docstring describing its purpose, parameters, and return value.

Add a descriptive docstring:

def load_sleap_config(cls, json_file_path: str) -> TrainerConfig:
+    """Load a training job configuration from a legacy SLEAP JSON config file.
+
+    Arguments:
+        json_file_path: Path to a legacy SLEAP JSON configuration file.
+
+    Returns:
+        A TrainingJobConfig instance containing the mapped configuration data.
+    """
    with open(json_file_path, "r") as f:
        old_config = json.load(f)

---

158-161: Missing error handling for mapper functions

There's no error handling for potential exceptions in the mapper functions.

Add error handling for the mapper functions:

-    data_config = data_mapper(old_config)
-    model_config = model_mapper(old_config)
-    trainer_config = trainer_mapper(old_config)
+    try:
+        data_config = data_mapper(old_config)
+        model_config = model_mapper(old_config)
+        trainer_config = trainer_mapper(old_config)
+    except Exception as e:
+        raise ValueError(f"Failed to map configuration from {json_file_path}: {str(e)}")

---

162-162: Filename attribute not set in the created config

Unlike the YAML loading methods, this function doesn't set the filename attribute to the source file path.

Set the filename attribute in the returned config:

-    return cls(data_config=data_config, model_config=model_config, trainer_config=trainer_config)
+    return cls(
+        data_config=data_config, 
+        model_config=model_config, 
+        trainer_config=trainer_config,
+        filename=json_file_path
+    )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (1)

189-243: ⚠️ Potential issue

Inconsistency in parameter handling

The function implementation has several issues that need to be addressed:

1. Many parameters are commented out (lines 191-201) including critical path parameters like train_labels_path and val_labels_path which are marked as MISSING in the DataConfig class.
2. Line 210 comments out use_augmentations_train, but it's used in the conditional at line 241.
3. The structure of the legacy_config appears inconsistent with some values coming from data.preprocessing and others from optimization.augmentation_config.

def data_mapper(legacy_config: dict) -> DataConfig:
    return DataConfig(
-        # train_labels_path=legacy_config.get("train_labels_path", MISSING),
-        # val_labels_path=legacy_config.get("val_labels_path", MISSING),
+        train_labels_path=legacy_config.get("train_labels_path", MISSING),
+        val_labels_path=legacy_config.get("val_labels_path", MISSING),
        # test_file_path=legacy_config.get("test_file_path"),
        # provider=legacy_config.get("provider", "LabelsReader"),
        # user_instances_only=legacy_config.get("user_instances_only", True),
        # data_pipeline_fw=legacy_config.get("data_pipeline_fw", "torch_dataset"),
        # np_chunks_path=legacy_config.get("np_chunks_path"),
        # litdata_chunks_path=legacy_config.get("litdata_chunks_path"),
        # use_existing_chunks=legacy_config.get("use_existing_chunks", False),
        # chunk_size=int(legacy_config.get("chunk_size", 100)),
        # delete_chunks_after_training=legacy_config.get("delete_chunks_after_training", True),
        preprocessing=PreprocessingConfig(
            is_rgb=legacy_config.get("data", {}).get("preprocessing", {}).get("ensure_rgb", False),
            max_height=legacy_config.get("data", {}).get("preprocessing", {}).get("target_height"),
            max_width=legacy_config.get("data", {}).get("preprocessing", {}).get("target_width"),
            scale=legacy_config.get("data", {}).get("preprocessing", {}).get("input_scaling", 1.0),
            crop_hw=legacy_config.get("data", {}).get("preprocessing", {}).get("crop_size"),
            min_crop_size=legacy_config.get("data", {}).get("preprocessing", {}).get("crop_size_detection_padding", 100),
        ),
-        # use_augmentations_train=legacy_config.get("use_augmentations_train", False),
+        use_augmentations_train=legacy_config.get("use_augmentations_train", False),
        augmentation_config=AugmentationConfig(
            intensity=IntensityConfig(
                uniform_noise_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("uniform_noise_min_val", 0.0),

sleap_nn/config/data_config.py (1)

189-190: Add docstring to data_mapper function

The data_mapper function lacks a docstring to explain its purpose, parameters, and return value. This is important for maintainability and understanding the code's intent.

+"""
+Convert a legacy SLEAP configuration dictionary to a DataConfig object.
+
+Args:
+    legacy_config: Dictionary containing configuration from older SLEAP version
+
+Returns:
+    DataConfig: A structured configuration object matching the legacy settings
+"""
def data_mapper(legacy_config: dict) -> DataConfig:

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/model_config.py (1)

906-917: ⚠️ Potential issue

Always creating centroid head violates one-of semantics.

Because centroid is set unconditionally, it will be active alongside any other head, contradicting the @oneof policy that only one head can be non-None. Wrap this in a condition to ensure that we only create the centroid head if head_type is set to "centroid".

-            centroid=CentroidConfig(
+            centroid=(
+                CentroidConfig(
+                    confmaps=CentroidConfMapsConfig(
+                        anchor_part=legacy_config.get("heads", {})
+                        .get("centroid", {})
+                        .get("anchor_part"),
+                        sigma=legacy_config.get("heads", {})
+                        .get("centroid", {})
+                        .get("sigma", 5.0),
+                        output_stride=legacy_config.get("heads", {})
+                        .get("centroid", {})
+                        .get("output_stride", 1),
+                    )
+                )
+                if legacy_config.get("head_type") == "centroid"
+                else None
+            ),

sleap_nn/config/data_config.py (1)

270-277: ⚠️ Potential issue

Scale set to (None, None) may break validation logic.

By default, GeometricConfig.scale expects numeric tuples (e.g., (0.9, 1.1)). Using (None, None) can trigger type errors or produce unintended outcomes. Provide valid numeric values if you want to enable scaling, or set this to None if scaling is not used.

-   scale=(
-       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_min", None),
-       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_max", None),
-   ),
+   scale=(
+       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_min", 0.9),
+       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_max", 1.1),
+   ),

sleap_nn/config/model_config.py (1)

835-836: Remove or clarify leftover comment.

These lines appear to contain a partial or uncertain comment about pretrained_backbone_weights (e.g., )?? # i think its different). This can be confusing to future maintainers. Consider removing or clarifying these comments to ensure the code is clean and unambiguous.

-        # pretrained_backbone_weights=legacy_config.get("PretrainedEncoderConfig")?? # i think its different
+        # pretrained_backbone_weights=legacy_config.get("PretrainedEncoderConfig")  # Clarify if needed or remove entirely

sleap_nn/config/data_config.py (1)

223-224: use_augmentations_train is commented out, causing mismatch.

The DataConfig instance sets use_augmentations_train to its default False, even if legacy_config indicates otherwise. Uncomment and map use_augmentations_train from legacy_config to keep the final config in sync.

-        # use_augmentations_train=legacy_config.get("use_augmentations_train", False),
+        use_augmentations_train=legacy_config.get("use_augmentations_train", False),

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/trainer_config.py (6)

236-236: Add a docstring to the trainer_mapper function.

The function is missing documentation explaining its purpose, parameters, and return value. A docstring would improve maintainability and help other developers understand the function's role.

 def trainer_mapper(legacy_config: dict) -> TrainerConfig:
+    """Convert a legacy configuration dictionary to a TrainerConfig instance.
+
+    This function extracts values from a legacy SLEAP configuration dictionary
+    and maps them to the appropriate fields in the TrainerConfig structure.
+
+    Args:
+        legacy_config: A dictionary containing legacy configuration parameters.
+
+    Returns:
+        A TrainerConfig instance populated with values from the legacy config.
+    """
     return TrainerConfig(

---

243-246: Decide on commented-out code.

There are numerous commented-out sections throughout the function. These make the code harder to read and maintain. Either implement these sections if they're needed or remove them if they're not.

If these configurations are planned for future implementation, consider adding a TODO comment indicating this, or remove them entirely if they're not needed.

---

311-350: Simplify conditional configuration creation.

The conditional creation of lr_scheduler is complex and nested. Consider extracting this logic to a separate helper function to improve readability.

+        lr_scheduler=_create_lr_scheduler_from_legacy(legacy_config),
-        lr_scheduler=(
-            LRSchedulerConfig(
-                # step_lr=StepLRConfig(
-                #     step_size=legacy_config.get("optimization", {})
-                #     .get("lr_scheduler", {})
-                #     .get("step_lr", {})
-                #     .get("step_size", 10),
-                #     gamma=legacy_config.get("optimization", {})
-                #     .get("lr_scheduler", {})
-                #     .get("step_lr", {})
-                #     .get("gamma", 0.1),
-                # ) if legacy_config.get("optimization", {}).get("lr_scheduler", {}).get("scheduler") == "StepLR" else None,
-                reduce_lr_on_plateau=ReduceLROnPlateauConfig(
-                    # threshold=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("threshold", 1e-4),
-                    # threshold_mode=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("threshold_mode", "rel"),
-                    # cooldown=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("cooldown", 0),
-                    patience=legacy_config.get("optimization", {})
-                    .get("learning_rate_schedule", {})
-                    .get("plateau_patience", 10),
-                    # factor=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("factor", 0.1),
-                    min_lr=legacy_config.get("optimization", {})
-                    .get("learning_rate_schedule", {})
-                    .get("min_learning_rate", 0.0),
-                )
-            )
-            if legacy_config.get("optimization", {}).get("learning_rate_schedule")
-            else None
-        ),

And add this helper function above:

def _create_lr_scheduler_from_legacy(legacy_config: dict) -> Optional[LRSchedulerConfig]:
    """Create LRSchedulerConfig from legacy configuration.
    
    Args:
        legacy_config: A dictionary containing legacy configuration parameters.
        
    Returns:
        A LRSchedulerConfig instance or None if learning_rate_schedule is not specified.
    """
    optimization = legacy_config.get("optimization", {})
    learning_rate_schedule = optimization.get("learning_rate_schedule")
    
    if not learning_rate_schedule:
        return None
        
    return LRSchedulerConfig(
        reduce_lr_on_plateau=ReduceLROnPlateauConfig(
            patience=learning_rate_schedule.get("plateau_patience", 10),
            min_lr=learning_rate_schedule.get("min_learning_rate", 0.0),
        )
    )

---

351-365: Similar to the lr_scheduler, extract early_stopping config creation to a helper function.

The conditional creation of early_stopping follows the same pattern as lr_scheduler and could benefit from the same refactoring approach.

+        early_stopping=_create_early_stopping_from_legacy(legacy_config),
-        early_stopping=(
-            EarlyStoppingConfig(
-                stop_training_on_plateau=legacy_config.get("optimization", {})
-                .get("learning_rate_schedule", {})
-                .get("reduce_on_plateau", False),
-                min_delta=legacy_config.get("optimization", {})
-                .get("learning_rate_schedule", {})
-                .get("plateau_min_delta", 0.0),
-                patience=legacy_config.get("optimization", {})
-                .get("learning_rate_schedule", {})
-                .get("plateau_patience", 1),
-            )
-            if legacy_config.get("optimization", {}).get("learning_rate_schedule")
-            else None
-        ),

And add this helper function below the _create_lr_scheduler_from_legacy function:

def _create_early_stopping_from_legacy(legacy_config: dict) -> Optional[EarlyStoppingConfig]:
    """Create EarlyStoppingConfig from legacy configuration.
    
    Args:
        legacy_config: A dictionary containing legacy configuration parameters.
        
    Returns:
        An EarlyStoppingConfig instance or None if learning_rate_schedule is not specified.
    """
    optimization = legacy_config.get("optimization", {})
    learning_rate_schedule = optimization.get("learning_rate_schedule")
    
    if not learning_rate_schedule:
        return None
        
    return EarlyStoppingConfig(
        stop_training_on_plateau=learning_rate_schedule.get("reduce_on_plateau", False),
        min_delta=learning_rate_schedule.get("plateau_min_delta", 0.0),
        patience=learning_rate_schedule.get("plateau_patience", 1),
    )

---

236-366: Add input validation for the legacy_config parameter.

The function doesn't validate the input legacy_config parameter. Consider adding validation to ensure it's a dictionary and contains the expected structure.

Add validation at the beginning of the function:

 def trainer_mapper(legacy_config: dict) -> TrainerConfig:
+    """Convert a legacy configuration dictionary to a TrainerConfig instance.
+
+    This function extracts values from a legacy SLEAP configuration dictionary
+    and maps them to the appropriate fields in the TrainerConfig structure.
+
+    Args:
+        legacy_config: A dictionary containing legacy configuration parameters.
+
+    Returns:
+        A TrainerConfig instance populated with values from the legacy config.
+
+    Raises:
+        TypeError: If legacy_config is not a dictionary.
+    """
+    if not isinstance(legacy_config, dict):
+        raise TypeError("legacy_config must be a dictionary")
+
     return TrainerConfig(

---

236-366: Consider using dictionary access helpers to simplify deeply nested dictionary accesses.

The function contains many instances of deeply nested dictionary access using chained .get() calls, which makes the code hard to read. Consider adding a helper function to simplify these accesses.

Add a helper function at the beginning of the file:

def get_nested_dict_value(dictionary: dict, keys: list, default=None):
    """Safely get a value from a nested dictionary.
    
    Args:
        dictionary: The dictionary to search in.
        keys: A list of keys representing the path to the value.
        default: The default value to return if the key path doesn't exist.
        
    Returns:
        The value at the key path, or the default value if not found.
    """
    current = dictionary
    for key in keys:
        if not isinstance(current, dict):
            return default
        current = current.get(key, {})
    return current or default

Then use it in the function:

-            batch_size=legacy_config.get("optimization", {}).get("batch_size", 1),
+            batch_size=get_nested_dict_value(legacy_config, ["optimization", "batch_size"], 1),

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/config/test_trainer_config.py (1)

239-297: Consider verifying the val_data_loader configuration as well.

The new test comprehensively checks most mapped parameters, but it might be beneficial to also verify that val_data_loader is assigned the expected defaults (e.g., batch_size, shuffle, etc.) since it mirrors the training data loader settings. This can increase coverage and potentially catch future regressions.

sleap_nn/config/trainer_config.py (1)

234-367: Remove or implement the commented-out fields for clarity.

Most fields are commented out (e.g., num_workers, amsgrad, step_lr) and can cause confusion about whether they are fully supported. Consider either removing them if they are no longer needed or fully implementing them. This will help maintain code clarity and consistency with the rest of the configuration.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/config/test_trainer_config.py (1)

239-297: Consider expanding test coverage for other checkpoint settings.
This test function looks comprehensive, verifying both mapped values and defaults. As an enhancement, consider adding an additional test scenario where "checkpointing" contains "latest_model": True" to confirm config.save_ckpt is correctly set to True.

+ legacy_config["optimization"]["checkpointing"] = {"latest_model": True}
+ config = trainer_mapper(legacy_config)
+ assert config.save_ckpt is True

tests/config/test_model_config.py (1)

123-171: Add optional scenarios for other heads.
The current test thoroughly verifies the single_instance head when using unet. A more exhaustive approach might also test other heads or backbone types if legacy configurations include them (e.g., centroid, centered_instance, multi_instance).

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (2)

277-284: 🛠️ Refactor suggestion

Fix scale parameter to use valid values instead of None.

The scale parameter is set with potentially None values, which could cause validation errors since the GeometricConfig class expects a list of floats with default (0.9, 1.1).

                    scale=(
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_min", None),
+                        .get("scale_min", 0.9),
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_max", None),
+                        .get("scale_max", 1.1),
                    ),

---

297-300: 🛠️ Refactor suggestion

Fix inconsistent handling of use_augmentations_train.

The function:

1. Has a commented-out conditional check for augmentation_config (lines 297-298)
2. Hardcodes use_augmentations_train=True (line 300)

This creates inconsistency with the DataConfig class default (False) and makes the augmentation configuration ignore the actual setting.

            )
-            # if legacy_config.get("use_augmentations_train", False)
-            # else None
+            if legacy_config.get("optimization", {}).get("use_augmentations", True)
+            else None
        ),
-        use_augmentations_train=True,
+        use_augmentations_train=legacy_config.get("optimization", {}).get("use_augmentations", True),

tests/assets/fixtures/datasets.py (1)

4-4: Remove the unused import to adhere to cleanliness.

Line 4 imports OmegaConf but it's never used in this file. Please remove it to keep the file lean and avoid unnecessary dependencies.

- from omegaconf import OmegaConf

🧰 Tools

🪛 Ruff (0.8.2)

4-4: omegaconf.OmegaConf imported but unused

Remove unused import: omegaconf.OmegaConf

(F401)

tests/config/test_training_job_config.py (3)

36-36: Remove the unused import for sleapnn_data_dir if it's not needed.

Static analysis indicates sleapnn_data_dir is unused. Review whether this fixture is needed; if not, remove it to eliminate the lint warning and keep dependencies minimal.

- from tests.assets.fixtures.datasets import sleapnn_data_dir, training_job_config_path
+ from tests.assets.fixtures.datasets import training_job_config_path

🧰 Tools

🪛 Ruff (0.8.2)

36-36: tests.assets.fixtures.datasets.sleapnn_data_dir imported but unused

Remove unused import

(F401)

---

36-36: tests.assets.fixtures.datasets.training_job_config_path imported but unused

Remove unused import

(F401)

---

38-38: Remove the unused import dataclasses.asdict.

This import does not appear to be used in the test file. Removing it helps tidy up the code.

- from dataclasses import asdict

🧰 Tools

🪛 Ruff (0.8.2)

38-38: dataclasses.asdict imported but unused

Remove unused import: dataclasses.asdict

(F401)

---

41-41: Remove the unused import json.

Although load_sleap_config relies on json, this test file itself does not call json directly. Removing the local import is recommended to satisfy static analysis warnings.

- import json

🧰 Tools

🪛 Ruff (0.8.2)

41-41: json imported but unused

Remove unused import: json

(F401)

sleap_nn/config/model_config.py (1)

837-1003: Consider broader support for non-UNet backbones.

Currently, this mapper only populates a UNetConfig and skips other potential backbone structures. If your legacy configs might include convnext or swint entries, consider handling them similarly. This ensures graceful fallback rather than silently ignoring those entries.

sleap_nn/config/trainer_config.py (3)

237-371: Add a docstring to the trainer_mapper function.

Adding a docstring would improve the maintainability and readability of this function by documenting its purpose, parameters, and return value.

+"""
+Maps a legacy configuration dictionary to a TrainerConfig object.
+
+Parameters:
+    legacy_config (dict): A dictionary containing legacy configuration settings.
+
+Returns:
+    TrainerConfig: An instance of TrainerConfig populated with values from legacy_config.
+"""
def trainer_mapper(legacy_config: dict) -> TrainerConfig:

---

305-309: Simplify the regex for capitalizing the first letter.

The regex to capitalize the first letter of the optimizer name is more complex than needed. Using the built-in capitalize() method would be more readable.

-        optimizer_name=re.sub(
-            r"^[a-z]",
-            lambda x: x.group().upper(),
-            legacy_config.get("optimization", {}).get("optimizer", "adam"),
-        ),
+        optimizer_name=legacy_config.get("optimization", {}).get("optimizer", "adam").capitalize(),

---

242-247: Explain the purpose of commented-out code sections.

There are multiple sections of commented-out code without explanation. This makes it difficult to understand whether these sections are incomplete implementations, deprecated code, or placeholders for future development.

Add explanatory comments at the beginning of commented-out sections, for example:

+        # NOTE: The following parameters are intentionally not mapped as they are not
+        # present in legacy SLEAP configurations or have different default values
         # num_workers=legacy_config.get("optimization", {})
         # .get("train_data_loader", {})
         # .get("num_workers", 0),

Also applies to: 250-255, 258-260, 262-268, 271-281, 282-304

sleap_nn/config/data_config.py (2)

190-191: Add a docstring to the data_mapper function.

The function lacks documentation explaining its purpose, parameters, and return value, which would improve maintainability.

+"""
+Maps a legacy configuration dictionary to a DataConfig object.
+
+Parameters:
+    legacy_config (dict): A dictionary containing legacy configuration settings.
+
+Returns:
+    DataConfig: An instance of DataConfig populated with values from legacy_config.
+"""
def data_mapper(legacy_config: dict) -> DataConfig:

---

236-239: Document the normalization of uniform_noise_max_val.

The function divides uniform_noise_max_val by 100.0, but this normalization is not documented and isn't applied to other similar values.

                    uniform_noise_max=legacy_config.get("optimization", {})
                    .get("augmentation_config", {})
-                    .get("uniform_noise_max_val", 100.0)
-                    / 100.0,
+                    .get("uniform_noise_max_val", 100.0) / 100.0,  # Convert percentage to fraction

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/config/test_training_job_config.py (1)

236-279: 🛠️ Refactor suggestion

Improve exception handling and file management in test.

The test has a few issues that should be addressed:

1. Unused variable e in the exception handler
2. Unsafe temporary file handling that doesn't guarantee cleanup
3. Complex control flow with try-except block

Consider using a context manager for temporary file handling and simplifying the test:

def test_load_sleap_config_from_file(training_job_config_path):
    """Test the load_sleap_config function with a sample legacy configuration from a JSON file."""
    # Path to the training_config.json file
    json_file_path = training_job_config_path

    # Load the configuration using the load_sleap_config method
    try:
        # Load the configuration using the load_sleap_config method
        config = load_sleap_config(TrainingJobConfig, json_file_path)
-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:
        with open(json_file_path, "r") as f:
            old_config = json.load(f)

        # Create a temporary file to hold the modified configuration
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.json', mode='w') as temp_file:
+        with tempfile.NamedTemporaryFile(suffix='.json', mode='w', delete=False) as temp_file:
            old_config['data']['labels']['training_labels'] = "notMISSING"
            old_config['data']['labels']['validation_labels'] = "notMISSING"
            
            json.dump(old_config, temp_file)
            temp_file_path = temp_file.name

-        config = load_sleap_config(TrainingJobConfig, temp_file_path)
-        os.remove(temp_file_path)
+        try:
+            config = load_sleap_config(TrainingJobConfig, temp_file_path)
+        finally:
+            os.remove(temp_file_path)

Also consider validating the exception message to ensure the correct failure is caught:

def test_load_sleap_config_missing_fields(training_job_config_path):
    """Test that load_sleap_config raises appropriate exceptions for missing fields."""
    with tempfile.NamedTemporaryFile(suffix='.json', mode='w') as temp_file:
        # Create a config with missing mandatory fields
        json.dump({"data": {"labels": {}}}, temp_file)
        temp_file.flush()
        
        # Verify the exception is raised with the expected message
        with pytest.raises(MissingMandatoryValue, match="train_labels_path is missing"):
            load_sleap_config(TrainingJobConfig, temp_file.name)

🧰 Tools

🪛 Ruff (0.8.2)

236-236: Redefinition of unused training_job_config_path from line 36

(F811)

---

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

tests/config/test_data_config.py (1)

127-200: Good test coverage for data_mapper function!

The test provides comprehensive validation for all aspects of the data_mapper function, including preprocessing, augmentation, and skeleton configurations.

Consider adding tests for error conditions:

1. What happens when mandatory fields like training_labels or validation_labels are missing?
2. How does the function handle edge cases or invalid values?

def test_data_mapper_missing_mandatory_fields():
    """Test the data_mapper function with missing mandatory fields."""
    # Test missing training_labels
    legacy_config = {
        "data": {
            "labels": {
                "validation_labels": "valid_path",
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="train_labels_path is missing"):
        data_mapper(legacy_config)
    
    # Test missing validation_labels
    legacy_config = {
        "data": {
            "labels": {
                "training_labels": "valid_path",
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="val_labels_path is missing"):
        data_mapper(legacy_config)

tests/config/test_training_job_config.py (3)

31-36: Clean up unused imports.

Several imports are not being used in this file. Consider removing them to improve code clarity.

from sleap_nn.config.training_job_config import TrainingJobConfig
from sleap_nn.config.training_job_config import load_sleap_config
from sleap_nn.config.model_config import ModelConfig
from sleap_nn.config.data_config import DataConfig
from sleap_nn.config.trainer_config import TrainerConfig, EarlyStoppingConfig
from sleap_nn.config.data_config import IntensityConfig
-from tests.assets.fixtures.datasets import sleapnn_data_dir, training_job_config_path
+from tests.assets.fixtures.datasets import training_job_config_path
from omegaconf import OmegaConf, MissingMandatoryValue
-from dataclasses import asdict
import json
-from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

36-36: tests.assets.fixtures.datasets.sleapnn_data_dir imported but unused

Remove unused import

(F401)

---

36-36: tests.assets.fixtures.datasets.training_job_config_path imported but unused

Remove unused import

(F401)

---

245-245: Remove unused exception variable.

The exception variable e is not being used in the exception handler.

    try:
        # Load the configuration using the load_sleap_config method
        config = load_sleap_config(TrainingJobConfig, json_file_path)
-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

251-260: Use context manager for safer temporary file handling.

The current approach doesn't guarantee the temporary file will be deleted if an exception occurs during loading.

-        with tempfile.NamedTemporaryFile(delete=False, suffix='.json', mode='w') as temp_file:
-            old_config['data']['labels']['training_labels'] = "notMISSING"
-            old_config['data']['labels']['validation_labels'] = "notMISSING"
-            
-            json.dump(old_config, temp_file)
-            temp_file_path = temp_file.name
-
-        config = load_sleap_config(TrainingJobConfig, temp_file_path)
-        os.remove(temp_file_path)
+        with tempfile.NamedTemporaryFile(suffix='.json', mode='w', delete=False) as temp_file:
+            old_config['data']['labels']['training_labels'] = "notMISSING"
+            old_config['data']['labels']['validation_labels'] = "notMISSING"
+            
+            json.dump(old_config, temp_file)
+            temp_file_path = temp_file.name
+
+        try:
+            config = load_sleap_config(TrainingJobConfig, temp_file_path)
+        finally:
+            if os.path.exists(temp_file_path):
+                os.remove(temp_file_path)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/model_config.py (1)

835-836: Remove unnecessary blank lines.

These two blank lines appear to serve no functional purpose. Consider removing them to keep the code more concise and aligned with standard style conventions.

tests/config/test_data_config.py (1)

127-200: Test provides comprehensive coverage of the data_mapper function.

The test effectively validates all aspects of the data_mapper function from sleap_nn/config/data_config.py, including:

1. Correct handling of labels paths
2. Preprocessing configuration conversion
3. Augmentation configuration handling
4. Intensity configuration parameters
5. Geometric transformation parameters
6. Skeletons configuration mapping

Consider adding test cases for error conditions, such as missing mandatory values (e.g., when train_labels_path or val_labels_path are missing). This would help ensure the function correctly raises appropriate exceptions in those scenarios.

def test_data_mapper_missing_mandatory_values():
    """Test that data_mapper raises appropriate exceptions for missing mandatory values."""
    # Test missing train_labels_path
    legacy_config_missing_train = {
        "data": {
            "labels": {
                "validation_labels": "some_path.slp"
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="train_labels_path is missing"):
        data_mapper(legacy_config_missing_train)
    
    # Test missing val_labels_path
    legacy_config_missing_val = {
        "data": {
            "labels": {
                "training_labels": "some_path.slp"
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="val_labels_path is missing"):
        data_mapper(legacy_config_missing_val)

tests/assets/fixtures/datasets.py (2)

4-4: Remove unused import.
The static analysis tool detects that OmegaConf is imported but unused.

Apply this diff to remove the unused import:

-from omegaconf import OmegaConf

🧰 Tools

🪛 Ruff (0.8.2)

4-4: omegaconf.OmegaConf imported but unused

Remove unused import: omegaconf.OmegaConf

(F401)

---

34-35: Update docstring to match file name.
The docstring says "Path to centered_instance_training_config file." but the actual file is "centered_instance_with_scaling_training_config.json".

-    """Path to centered_instance_training_config file."""
+    """Path to centered_instance_with_scaling_training_config file."""

tests/config/test_training_job_config.py (8)

36-36: Avoid wildcard import.
Using import * can mask undefined names.

Here is a diff showing an explicit import:

-from tests.assets.fixtures.datasets import *
+from tests.assets.fixtures.datasets import (
+    sleapnn_data_dir,
+    training_job_config_path,
+    bottomup_training_config_path,
+    centered_instance_training_config_path,
+    centered_instance_with_scaling_training_config_path,
+    centroid_training_config_path,
+    single_instance_training_config_path,
+    topdown_training_config_path,
+)

🧰 Tools

🪛 Ruff (0.8.2)

36-36: from tests.assets.fixtures.datasets import * used; unable to detect undefined names

(F403)

---

38-38: Remove unused import.
dataclasses.asdict is not referenced in this file.

-from dataclasses import asdict

🧰 Tools

🪛 Ruff (0.8.2)

38-38: dataclasses.asdict imported but unused

Remove unused import: dataclasses.asdict

(F401)

---

42-42: Remove unused import.
omegaconf.MISSING is never used.

-from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

42-42: omegaconf.MISSING imported but unused

Remove unused import: omegaconf.MISSING

(F401)

---

43-43: Remove unused import.
pprint.pprint is never used.

-from pprint import pprint

🧰 Tools

🪛 Ruff (0.8.2)

43-43: pprint.pprint imported but unused

Remove unused import: pprint.pprint

(F401)

---

245-245: Remove assignment to unused variable.
The local variable e is assigned but never used.

-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

296-296: Remove assignment to unused variable.
The local variable e is assigned but never used.

-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

296-296: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

340-340: Remove assignment to unused variable.
Again, e is not utilized.

-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

340-340: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

237-285: Refactor repeated code blocks in tests.
Each test repeats the same logic for handling MissingMandatoryValue, rewriting JSON, and loading the updated file.

Consider extracting this into a helper function or fixture to reduce duplication:

def load_or_fix_legacy_config(json_file_path, required_keys=None):
    """
    Attempt to load the config. If MissingMandatoryValue is raised,
    rewrite the config JSON with non-missing placeholders.
    """
    from omegaconf import MissingMandatoryValue
    import tempfile, os, json

    if required_keys is None:
        required_keys = [("data", "labels", "training_labels"), ("data", "labels", "validation_labels")]
    try:
        return load_sleap_config(TrainingJobConfig, json_file_path)
    except MissingMandatoryValue:
        with open(json_file_path, "r") as f:
            old_config = json.load(f)

        for key_path in required_keys:
            cfg = old_config
            for kp in key_path[:-1]:
                cfg = cfg[kp]
            cfg[key_path[-1]] = "notMISSING"

        with tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode="w") as temp_file:
            json.dump(old_config, temp_file)
            temp_file_path = temp_file.name

        config = load_sleap_config(TrainingJobConfig, temp_file_path)
        os.remove(temp_file_path)
        return config

Also applies to: 287-327, 329-371, 373-407, 409-439, 441-473, 475-510

🧰 Tools

🪛 Ruff (0.8.2)

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (1)

190-313: 🛠️ Refactor suggestion

Several concerns in the new data_mapper function.

1. Commented-out code: Key parameters like provider, user_instances_only, etc., remain commented out. If they are required, consider uncommenting or removing them altogether to reduce confusion.
2. Forcing use_augmentations_train to True: You currently override this with a hard-coded True (line 310). Most designs expect reading it from the legacy config (e.g., legacy_config.get("use_augmentations_train", False)).
3. Potential (None, None) usage: Passing (None, None) as scale (lines 288–293) can break validations in GeometricConfig. Confirm whether this is intended or provide a default numeric range if the config is missing a scale setting.

Below is a snippet to let legacy_config decide if augmentations are used, and to supply a safer default for scale:

-        use_augmentations_train=True,
+        use_augmentations_train=bool(
+            legacy_config.get("optimization", {})
+                       .get("use_augmentations", False)
+        ),

-                    scale=(
-                        legacy_config.get("optimization", {})
-                        .get("augmentation_config", {})
-                        .get("scale_min", None),
-                        legacy_config.get("optimization", {})
-                        .get("augmentation_config", {})
-                        .get("scale_max", None),
-                    ),
+                    scale=(
+                        legacy_config.get("optimization", {})
+                        .get("augmentation_config", {})
+                        .get("scale_min", 0.9),
+                        legacy_config.get("optimization", {})
+                        .get("augmentation_config", {})
+                        .get("scale_max", 1.1),
+                    ),

sleap_nn/config/training_job_config.py (1)

155-173: 🛠️ Refactor suggestion

Multiple issues with the load_sleap_config function.

1. Lack of error handling: No try/except to handle file I/O errors and JSON parse errors (e.g., FileNotFoundError, json.JSONDecodeError). This can lead to uncaught exceptions at runtime.
2. @classmethod usage: The function uses cls as its first parameter, but it isn’t decorated with @classmethod, causing potential confusion about its intended usage.
3. Missing check_output_strides: Unlike from_yaml and load_yaml, this loader doesn’t invoke TrainingJobConfig.check_output_strides(config), potentially leading to inconsistent stride settings relative to the YAML loading path.
4. Return type mismatch: The function signature says it returns a TrainerConfig, but the actual return object is an OmegaConf instance.

Below is a suggested refactor that includes error handling, calls check_output_strides, and aligns the return type:

+@classmethod
-def load_sleap_config(cls, json_file_path: str) -> TrainerConfig:
+def load_sleap_config(cls, json_file_path: str) -> OmegaConf:
    +import json
    try:
        with open(json_file_path, "r") as f:
            old_config = json.load(f)
    except FileNotFoundError:
        raise FileNotFoundError(f"Configuration file not found: {json_file_path}")
    except json.JSONDecodeError:
        raise ValueError(f"Invalid JSON in configuration file: {json_file_path}")

    data_config = data_mapper(old_config)
    model_config = model_mapper(old_config)
    trainer_config = trainer_mapper(old_config)

    config = cls(
        data_config=data_config,
        model_config=model_config,
        trainer_config=trainer_config,
    )

    schema = OmegaConf.structured(config)
    config_omegaconf = OmegaConf.merge(schema, OmegaConf.create(asdict(config)))
    +config_omegaconf = cls.check_output_strides(config_omegaconf)
    OmegaConf.to_container(config_omegaconf, resolve=True, throw_on_missing=True)

    return config_omegaconf

sleap_nn/config/data_config.py (1)

8-8: Remove the unused import.
MissingMandatoryValue is never invoked, so please remove it to address the linter’s warning and keep your imports clean.

-from omegaconf import MISSING, MissingMandatoryValue
+from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

8-8: omegaconf.MissingMandatoryValue imported but unused

Remove unused import: omegaconf.MissingMandatoryValue

(F401)

tests/config/test_training_job_config.py (2)

36-37: Avoid star-import if possible.
Using import * from fixtures can clutter the namespace and make it harder to track dependencies. Prefer explicitly importing only what you need.

🧰 Tools

🪛 Ruff (0.8.2)

36-36: from tests.assets.fixtures.datasets import * used; unable to detect undefined names

(F403)

---

41-43: Several unused imports.

• json is used in the tests below, so it’s fine.
• MISSING (line 42) and pprint (line 43) appear unused. Consider removing them along with any other imports flagged by static analysis for clarity.

 import json
-from omegaconf import MISSING
-from pprint import pprint

🧰 Tools

🪛 Ruff (0.8.2)

42-42: omegaconf.MISSING imported but unused

Remove unused import: omegaconf.MISSING

(F401)

---

43-43: pprint.pprint imported but unused

Remove unused import: pprint.pprint

(F401)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (2)

295-302: 🛠️ Refactor suggestion

Fix scale parameter default values

The scale parameter in GeometricConfig is set to None values, which could conflict with validators that expect numeric values. In the GeometricConfig class definition (line 113), the default is (0.9, 1.1).

                    scale=(
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_min", None),
+                        .get("scale_min", 0.9),
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_max", None),
+                        .get("scale_max", 1.1),
                    ),

---

315-318: ⚠️ Potential issue

Fix inconsistent use_augmentations_train handling

There are two issues:

1. use_augmentations_train is hardcoded to True regardless of what's in the legacy configuration
2. The conditional check for setting augmentation_config to None is commented out

This could lead to inconsistencies between the augmentation configuration and whether augmentations are actually used.

            )
-            # if legacy_config.get("use_augmentations_train", False)
-            # else None
+            if legacy_config.get("optimization", {}).get("use_augmentations", True)
+            else None
        ),
-        use_augmentations_train=True,
+        use_augmentations_train=legacy_config.get("optimization", {}).get("use_augmentations", True),

sleap_nn/config/trainer_config.py (4)

313-317: Consider simplifying the regex for capitalizing the optimizer name.

The regex approach works but could be simplified using standard string methods.

-    optimizer_name=re.sub(
-        r"^[a-z]",
-        lambda x: x.group().upper(),
-        legacy_config.get("optimization", {}).get("optimizer", "adam"),
-    ),
+    optimizer_name=legacy_config.get("optimization", {}).get("optimizer", "adam").capitalize(),

---

270-311: Consider adding a TODO comment for commented-out configurations.

There are numerous commented-out configurations in this function. If these are planned for future implementation, consider adding a TODO comment explaining the plan. If they're not needed, consider removing them for better code clarity.

---

324-363: Consider adding a helper function for nested dictionary access.

The deeply nested get() calls reduce readability. Consider extracting a helper function to make the code more maintainable.

def get_nested(d, *keys, default=None):
    """Get a value from a nested dictionary by a sequence of keys."""
    result = d
    for key in keys:
        if not isinstance(result, dict):
            return default
        result = result.get(key, {})
    return result if result != {} else default

# Usage example:
patience = get_nested(legacy_config, "optimization", "learning_rate_schedule", "plateau_patience", default=10)

---

349-358: Inconsistent path access could lead to confusion.

Some parameters are accessed from "learning_rate_schedule" while others might come from "lr_scheduler" (in commented code). This inconsistency could lead to confusion.

Consider documenting why different paths are used or standardize on a consistent approach if possible.

sleap_nn/config/model_config.py (4)

846-895: Consider simplifying repetitive dictionary access patterns

The code contains many repetitive nested .get() calls which make the code verbose and harder to maintain. Consider creating a helper function to simplify this pattern.

+def _get_nested(config, path, default=None):
+    """Get a value from a nested dictionary using a list of keys.
+    
+    Args:
+        config: Dictionary to access
+        path: List of keys to traverse
+        default: Default value if path doesn't exist
+        
+    Returns:
+        Value at the path or default if not found
+    """
+    result = config
+    for key in path:
+        if not isinstance(result, dict):
+            return default
+        result = result.get(key, {})
+    return result if result != {} else default

 def model_mapper(legacy_config: dict) -> ModelConfig:
     """Map the legacy model configuration to the new model configuration.
 
     Args:
         legacy_config: A dictionary containing the legacy model configuration.
 
     Returns:
         An instance of `ModelConfig` with the mapped configuration.
     """
+    model_config = legacy_config.get("model", {})
     return ModelConfig(
         # init_weights=legacy_config.get("init_weights", "default"),
         # pre_trained_weights not in old config
         # pretrained_backbone_weights=legacy_config.get("PretrainedEncoderConfig")?? # i think its different
         # pretrained_head_weights not in old config
         backbone_config=BackboneConfig(
             unet=(
                 UNetConfig(
                     # in_channels=legacy_config.get("backbone", {}).get("in_channels", 1),
                     # kernel_size=legacy_config.get("backbone", {}).get("kernel_size", 3),
-                    filters=legacy_config.get("model", {})
-                    .get("backbone", {})
-                    .get("unet", {})
-                    .get("filters", 32),
+                    filters=_get_nested(model_config, ["backbone", "unet", "filters"], 32),

This pattern could be applied throughout the function to make it more maintainable.

---

896-920: Remove unnecessary nested parentheses

The SingleInstanceConfig has an unnecessary level of parentheses in the conditional, which can be simplified.

         head_configs=HeadConfig(
             single_instance=(
-                (
-                    SingleInstanceConfig(
-                        confmaps=SingleInstanceConfMapsConfig(
-                            part_names=legacy_config.get("model", {})
-                            .get("heads", {})
-                            .get("single_instance", {})
-                            .get("part_names"),
-                            sigma=legacy_config.get("model", {})
-                            .get("heads", {})
-                            .get("single_instance", {})
-                            .get("sigma", 5.0),
-                            output_stride=legacy_config.get("model", {})
-                            .get("heads", {})
-                            .get("single_instance", {})
-                            .get("output_stride", 1),
-                        )
-                    )
-                )
+                SingleInstanceConfig(
+                    confmaps=SingleInstanceConfMapsConfig(
+                        part_names=legacy_config.get("model", {})
+                        .get("heads", {})
+                        .get("single_instance", {})
+                        .get("part_names"),
+                        sigma=legacy_config.get("model", {})
+                        .get("heads", {})
+                        .get("single_instance", {})
+                        .get("sigma", 5.0),
+                        output_stride=legacy_config.get("model", {})
+                        .get("heads", {})
+                        .get("single_instance", {})
+                        .get("output_stride", 1),
+                    )
+                )
                 if legacy_config.get("model", {})
                 .get("heads", {})
                 .get("single_instance")
                 else None
             ),

---

969-1016: The bottomup head mapping uses multi_instance paths

Good job on adding the conditional check to maintain the @oneof constraint for the bottomup head. I noticed that the legacy configuration key is multi_instance while the new configuration key is bottomup. This seems intentional, but it would be helpful to add a comment explaining this mapping for future maintenance.

             bottomup=(
                 BottomUpConfig(
+                    # Note: bottomup head was previously called multi_instance in legacy configs
                     confmaps=BottomUpConfMapsConfig(
                         loss_weight=legacy_config.get("model", {})
                         .get("heads", {})
                         .get("multi_instance", {})

---

837-1018: Add input validation for legacy_config

The function assumes that legacy_config is a valid dictionary, but doesn't validate this. Consider adding input validation to handle edge cases gracefully.

 def model_mapper(legacy_config: dict) -> ModelConfig:
     """Map the legacy model configuration to the new model configuration.
 
     Args:
         legacy_config: A dictionary containing the legacy model configuration.
 
     Returns:
         An instance of `ModelConfig` with the mapped configuration.
     """
+    if not isinstance(legacy_config, dict):
+        logger.warning("Invalid legacy_config provided, using defaults")
+        legacy_config = {}
+    
     return ModelConfig(

sleap_nn/config/data_config.py (4)

8-8: Remove unused import

The MissingMandatoryValue exception is imported but never used in the code.

-from omegaconf import MISSING, MissingMandatoryValue
+from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

8-8: omegaconf.MissingMandatoryValue imported but unused

Remove unused import: omegaconf.MissingMandatoryValue

(F401)

---

209-216: Add explanation for commented-out parameters

Several parameters are commented out without explanation. If these parameters aren't used in legacy configurations, add a comment explaining why they're commented out. If they should be included, uncomment them.

-        # provider=legacy_config.get("provider", "LabelsReader"),
-        # user_instances_only=legacy_config.get("user_instances_only", True),
-        # data_pipeline_fw=legacy_config.get("data_pipeline_fw", "torch_dataset"),
-        # np_chunks_path=legacy_config.get("np_chunks_path"),
-        # litdata_chunks_path=legacy_config.get("litdata_chunks_path"),
-        # use_existing_chunks=legacy_config.get("use_existing_chunks", False),
-        # chunk_size=int(legacy_config.get("chunk_size", 100)),
-        # delete_chunks_after_training=legacy_config.get("delete_chunks_after_training", True),
+        # These parameters are not part of the legacy SLEAP configuration format
+        # provider=legacy_config.get("provider", "LabelsReader"),
+        # user_instances_only=legacy_config.get("user_instances_only", True),
+        # data_pipeline_fw=legacy_config.get("data_pipeline_fw", "torch_dataset"),
+        # np_chunks_path=legacy_config.get("np_chunks_path"),
+        # litdata_chunks_path=legacy_config.get("litdata_chunks_path"),
+        # use_existing_chunks=legacy_config.get("use_existing_chunks", False),
+        # chunk_size=int(legacy_config.get("chunk_size", 100)),
+        # delete_chunks_after_training=legacy_config.get("delete_chunks_after_training", True),

---

303-312: Explain commented-out geometric augmentation parameters

Several geometric augmentation parameters are commented out without explanation. If these parameters aren't available in legacy configurations, add a comment explaining why they're commented out. If they should be included, uncomment them.

-                    # translate_width=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_width", 0.2),
-                    # translate_height=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_height", 0.2),
-                    # affine_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("affine_p", 0.0),
-                    # erase_scale_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_min", 0.0001),
-                    # erase_scale_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_max", 0.01),
-                    # erase_ratio_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_min", 1.0),
-                    # erase_ratio_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_max", 1.0),
-                    # erase_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_p", 0.0),
-                    # mixup_lambda=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_lambda", [0.01, 0.05]),
-                    # mixup_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_p", 0.0),
+                    # These geometric augmentation parameters are not available in legacy SLEAP configuration format
+                    # translate_width=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_width", 0.2),
+                    # translate_height=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_height", 0.2),
+                    # affine_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("affine_p", 0.0),
+                    # erase_scale_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_min", 0.0001),
+                    # erase_scale_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_max", 0.01),
+                    # erase_ratio_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_min", 1.0),
+                    # erase_ratio_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_max", 1.0),
+                    # erase_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_p", 0.0),
+                    # mixup_lambda=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_lambda", [0.01, 0.05]),
+                    # mixup_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_p", 0.0),

---

318-318: Add unit tests for data_mapper function

The data_mapper function is complex and handles critical configuration mapping, but there don't appear to be tests for it. Consider adding unit tests to verify it handles various legacy configuration formats correctly.

Would you like me to generate example unit tests for the data_mapper function?